# ==============================================================================
# 文件: feature_extractor.py
# 作用: 负责加载Hugging Face的预训练模型（如CLIP, SigLIP），并使用它
#      为输入的数据集提取嵌入向量（embeddings）。
# ==============================================================================
import torch
from transformers import AutoProcessor, AutoModel # 从Hugging Face的transformers库导入模型和处理器
from torch.utils.data import DataLoader
from tqdm import tqdm # 导入tqdm库，用于显示进度条
import os # 导入os库，用于文件操作

import config # 从同级目录的config模块导入配置

def get_model_and_processor():
    """
    根据config中的模型名称，从Hugging Face加载预训练模型和对应的预处理器。
    :return: model, processor
    """
    print(f"正在从Hugging Face加载模型: {config.FEATURE_EXTRACTOR_NAME}...")
    # 加载预处理器，它负责将图像转换为模型所需的格式
    processor = AutoProcessor.from_pretrained(config.FEATURE_EXTRACTOR_NAME)
    # 加载预训练模型
    model = AutoModel.from_pretrained(config.FEATURE_EXTRACTOR_NAME).to(config.DEVICE)
    # 将模型设置为评估模式，这会关闭Dropout等只在训练时使用的层
    model.eval()
    return model, processor

def extract_embeddings(dataset, model, processor, description="Extracting"):
    """
    为给定的数据集提取嵌入向量。
    :param dataset: PyTorch Dataset对象。
    :param model: 加载好的预训练模型。
    :param processor: 加载好的预处理器。
    :param description: str, tqdm进度条的描述文字。
    :return: 包含所有嵌入向量的Tensor。
    """
    # 创建一个数据加载器，用于批量处理数据
    dataloader = DataLoader(dataset, batch_size=128, shuffle=False)
    
    all_embeddings = [] # 创建一个空列表来存放所有批次的嵌入

    # 使用torch.no_grad()上下文管理器，临时禁用梯度计算，以加速计算并节省内存
    with torch.no_grad():
        # 使用tqdm显示进度条
        for batch in tqdm(dataloader, desc=description):
            # batch[0] 是图像张量
            images = batch[0].to(config.DEVICE)
            
            # 使用预处理器处理图像
            # `do_rescale=False` 是因为torchvision的ToTensor已经将像素值缩放到[0,1]
            inputs = processor(images=images, return_tensors="pt", do_rescale=False).to(config.DEVICE)
            
            # 从模型中获取图像特征（嵌入）
            # 对于CLIP和SigLIP，我们通常使用image_embeds
            image_features = model.get_image_features(**inputs)
            
            # 将嵌入移动到CPU并添加到列表中
            all_embeddings.append(image_features.cpu())

    # 将所有批次的嵌入拼接成一个大的Tensor
    return torch.cat(all_embeddings, dim=0)

def get_or_create_embeddings(p_dataset, u_dataset):
    """
    检查嵌入文件是否存在。如果存在，则直接加载；如果不存在，则提取、保存并返回。
    :param p_dataset: P集 Dataset对象。
    :param u_dataset: U集 Dataset对象。
    :return: p_embeddings (Tensor), u_embeddings (Tensor)
    """
    # 检查配置文件中指定的嵌入文件是否存在
    if os.path.exists(config.EMBEDDING_FILE_PATH):
        print(f"发现已存在的嵌入文件，正在加载: {config.EMBEDDING_FILE_PATH}")
        # 加载保存的嵌入数据
        data = torch.load(config.EMBEDDING_FILE_PATH)
        return data['p_embeddings'], data['u_embeddings']
    else:
        print("未发现嵌入文件，开始提取...")
        # 加载模型和处理器
        model, processor = get_model_and_processor()
        
        # 分别为P集和U集提取嵌入
        p_embeddings = extract_embeddings(p_dataset, model, processor, "提取P集嵌入")
        u_embeddings = extract_embeddings(u_dataset, model, processor, "提取U集嵌入")
        
        print(f"嵌入提取完成，正在保存到: {config.EMBEDDING_FILE_PATH}")
        # 将提取到的嵌入以字典形式保存到文件
        torch.save({
            'p_embeddings': p_embeddings,
            'u_embeddings': u_embeddings
        }, config.EMBEDDING_FILE_PATH)
        
        return p_embeddings, u_embeddings
